package org.example;

import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.IpSelector;
import org.pcap4j.packet.UdpPacket;

import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class ServerMain {


    private static ConcurrentHashMap<String,AsynchronousSocketChannel> channelMap = new ConcurrentHashMap<>();


    // 创建线程池，使用虚拟线程
    private static final ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor();

    private static IPPool ipPool = new IPPool();


    public static void main(String[] args) {
        // 设置端口号
        int port = Integer.parseInt(Optional.ofNullable(System.getenv("PORT")).orElse("8888"));
        // 创建异步服务器套接字通道
        try (AsynchronousServerSocketChannel serverChannel = AsynchronousServerSocketChannel.open()) {
            // 绑定到本地端口
            serverChannel.bind(new InetSocketAddress(port));

            System.out.println("服务器启动，监听端口 " + port);

            // 接受连接
            while (true) {
                // 接受客户端连接
                AsynchronousSocketChannel clientChannel = serverChannel.accept().get();
                System.out.println("新连接： " + clientChannel.getRemoteAddress());

                // 使用虚拟线程处理每个连接
                executorService.submit(() -> handleClient(clientChannel));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static void handleClient(AsynchronousSocketChannel clientChannel) {
        try {

            pushChannelMap(clientChannel);

            var buff = ByteBuffer.allocate(4);
            while (!Thread.interrupted()){
                var len = clientChannel.read(buff).get();
                if(len <= 0){
                    break;
                }
                buff.flip();
                var anInt = buff.getInt();
                if (anInt<=0 || anInt > 1024 * 1024) {
                    System.out.println("无效的包长度: " + anInt);
                    break;
                }
                var buff_dy = ByteBuffer.allocate(anInt);
                var len_dy = clientChannel.read(buff_dy).get();
                if(len_dy <= 0){
                    break;
                }
                var bt = new byte[len_dy];
                buff_dy.flip();
                buff_dy.get(bt);
                parsePacket(bt);
                buff.clear();
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                clientChannel.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static void pushChannelMap(AsynchronousSocketChannel clientChannel) throws Exception {
        var ip = ipPool.acquireIP();
        if(ip == null){
            throw new RuntimeException("IP池已满，无法分配新的IP");
        }

        byte[] ipBytes = ip.getBytes();
        var buffer = ByteBuffer.allocate(4+ipBytes.length);
        buffer.putInt(ipBytes.length);
        buffer.put(ipBytes);
        buffer.flip();

        //写入客户端
        clientChannel.write(buffer).get();

        channelMap.put(ip,clientChannel);
    }

    private static void parsePacket(byte[] data){

        try {
            IpPacket packet = (IpPacket) IpSelector.newPacket(data, 0, data.length);
            InetAddress dstAddr = packet.getHeader().getDstAddr();
            InetAddress srcAddr = packet.getHeader().getSrcAddr();


            if(dstAddr instanceof Inet4Address){
                Inet4Address dst4Addr = (Inet4Address)dstAddr;
                Inet4Address src4Addr = (Inet4Address)srcAddr;


                String hostAddress = src4Addr.getHostAddress();

                ByteBuffer allocate = ByteBuffer.allocate(data.length + 4);
                allocate.putInt(data.length);
                allocate.put(data);
                allocate.flip();

                if(dst4Addr.getAddress()[3] == (byte) 255){
                    //广播发送给全部channel
                    for (var entry : channelMap.entrySet()) {
                        var key = entry.getKey();
                        var value = entry.getValue();

                        if(key.equals(hostAddress)){
                            //不需要广播自己
                            continue;
                        }

                        try {
                            value.write(allocate).get();
                        }catch (Exception e){
                            e.printStackTrace();
                            value.close();
                            channelMap.remove(key);
                            ipPool.releaseIP(key);
                        }

                    }

                }else{
                    String key = dst4Addr.getHostAddress();
                    var asc = channelMap.get(key);
                    if(asc!=null){
                        try {
                            asc.write(allocate).get();
                        }catch (Exception e){
                            e.printStackTrace();
                            asc.close();
                            channelMap.remove(key);
                            ipPool.releaseIP(key);
                        }
                    }

                }
            }


            if(packet.getPayload() instanceof UdpPacket){

            }


        }catch (Exception e){
            e.printStackTrace();
        }
    }

}